import datetime
import os
import shutil


# 生成Yolo VOC xml格式
class GenerateJpgAndXml:
    """
    参数名含义：
    parentName:存放jpg和xml上一级文件夹名字,如person

    """

    def __init__(self, parentName, labelDict):
        self.parentName = parentName
        # 存放所有文件的主文件夹路径
        self.parentPath = r"G:\workspace\python\yolov8-all-in-one\datasets\VOCdevkit"
        self.xmlPath = os.path.join(self.parentPath, "Annotations")
        self.imgPath = os.path.join(self.parentPath, 'JPEGImages')
        # 存放标签的字典
        self.labelDict = labelDict
        # 第一次进来，需要判断下文件夹是否存在
        self.isExist()

    def isExist(self):
        # 存放xml文件的文件夹
        if not os.path.exists(self.xmlPath):
            os.makedirs(self.xmlPath)
        if not os.path.exists(self.imgPath):
            os.makedirs(self.imgPath)

    def generatr_xml(self, jpgPath, objName, xmin, ymin, xmax, ymax):
        # 获取当前时间戳
        xmlPrefix = datetime.datetime.now().strftime("%Y%m%d%H%M%S%f")
        # print(xmlPrefix)
        # jpg名字
        jpgName = xmlPrefix + ".jpg"
        # jpg路径
        imgPath = os.path.join(self.imgPath, jpgName)
        shutil.copyfile(jpgPath, imgPath)
        # 写图片
        # xml路径
        xmlPath = os.path.join(self.xmlPath, xmlPrefix + ".xml")
        with open(xmlPath, 'w', encoding="utf-8") as xml_file:
            xml_file.write('<annotation>\n')
            xml_file.write('\t<folder>' + self.parentName + '</folder>\n')
            xml_file.write('\t<filename>' + jpgName + '</filename>\n')
            xml_file.write('\t<path>' + imgPath + '</path>\n')
            xml_file.write('\t<source>\n')
            xml_file.write('\t\t<database>' + 'Unknown' + '</database>\n')
            xml_file.write('\t</source>\n')
            xml_file.write('\t<size>\n')
            xml_file.write('\t\t<width>' + str(xmax) + '</width>\n')
            xml_file.write('\t\t<height>' + str(ymax) + '</height>\n')
            xml_file.write('\t\t<depth>' + str(3) + '</depth>\n')
            xml_file.write('\t</size>\n')
            xml_file.write('\t<segmented>0</segmented>\n')

            xml_file.write('\t<object>\n')
            xml_file.write('\t\t<name>' + objName + '</name>\n')
            xml_file.write('\t\t<pose>Unspecified</pose>\n')
            xml_file.write('\t\t<truncated>0</truncated>\n')
            xml_file.write('\t\t<difficult>0</difficult>\n')
            xml_file.write('\t\t<bndbox>\n')
            xml_file.write('\t\t\t<xmin>' + str(xmin) + '</xmin>\n')
            xml_file.write('\t\t\t<ymin>' + str(ymin) + '</ymin>\n')
            xml_file.write('\t\t\t<xmax>' + str(xmax) + '</xmax>\n')
            xml_file.write('\t\t\t<ymax>' + str(ymax) + '</ymax>\n')
            # xml_file.write('\t\t\t<angle>' + str(4) + '</angle>\n')
            xml_file.write('\t\t</bndbox>\n')
            # xml_file.write('\t\t<extra/>\n')
            xml_file.write('\t</object>\n')
            xml_file.write('</annotation>')
        # customPrint(f"{jpgPath}的jpg和xml已写入")
